import socket
import json
import time
from config import Config
import logging

logger = logging.getLogger(__name__)

class TCPClient:
    def __init__(self, host=Config.TCP_CLIENT_HOST, port=Config.TCP_PORT):
        self.host = host
        self.port = port
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.connect()
        
    def connect(self):
        """连接到服务器"""
        retry_count = 0
        max_retries = 3
        retry_delay = 2  # 秒
        
        while retry_count < max_retries:
            try:
                self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                logger.info(f"正在连接服务器 {self.host}:{self.port}")
                self.sock.connect((self.host, self.port))
                logger.info("连接成功")
                return True
            except ConnectionRefusedError:
                retry_count += 1
                if retry_count < max_retries:
                    logger.warning(f"连接失败，{retry_delay}秒后重试 ({retry_count}/{max_retries})")
                    time.sleep(retry_delay)
                else:
                    logger.error("连接服务器失败，已达到最大重试次数")
                    raise
            except Exception as e:
                logger.error(f"连接服务器时发生错误: {str(e)}")
                raise
            
    def send_request(self, request):
        """发送请求并接收响应"""
        try:
            self.sock.sendall(json.dumps(request).encode())
            response = self.sock.recv(Config.TCP_BUFFER_SIZE).decode()
            return json.loads(response)
        except Exception as e:
            print(f"请求失败: {str(e)}")
            return {'error': str(e)}
            
    def add_documents(self, url: str):
        """通过URL添加文档"""
        request = {
            'type': 'add_documents',
            'data': url  # 发送文档URL而不是目录路径
        }
        return self.send_request(request)
        
    def search(self, query: str, top_k: int = Config.DEFAULT_TOP_K):
        """检索文档"""
        request = {
            'type': 'search',
            'data': query,
            'top_k': 1
        }
        return self.send_request(request)
        
    def get_documents(self, indices: list):
        """获取文档内容"""
        request = {
            'type': 'get_documents',
            'indices': indices
        }
        return self.send_request(request)
        
    def close(self):
        """关闭连接"""
        self.sock.close() 